import pathlib
import typing

import matplotlib.pyplot
import torch

import models
import shared.checkpoint
import shared.torchtools


def plot_checkpoint(checkpoint: pathlib.Path,
                    times: typing.Sequence[float] = (0.20, 0.80)) \
        -> matplotlib.pyplot.Figure:
    states = torch.load(checkpoint)

    model = models.PinnOfKortewegDeVriesEquation()
    model.load_state_dict(states["model"])
    model.eval()

    print(f"lambda1: {float(model.lambda1)}")
    print(f"lambda2: {float(model.lambda2)}")

    figure = matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    def _plot_contour():
        axes = figure.add_subplot(2, 1, 1)

        t = torch.linspace(0, 1, 1000)
        x = torch.linspace(-1, 1, 1000)
        t, x = torch.meshgrid(t, x, indexing="xy")

        u = model(torch.flatten(t), torch.flatten(x))
        u = shared.torchtools.reshape_as(u, t, detach=True)

        contour = axes.contourf(t, x, u,
                                levels=80, cmap=matplotlib.pyplot.get_cmap("rainbow"))

        axes.set_title("$u(t, x)$")
        axes.set_xlabel("$t$")
        axes.set_ylabel("$x$")
        figure.colorbar(contour, fraction=0.08, pad=0.01)

    _plot_contour()

    def _plot_times():
        for i, time in enumerate(times, start=len(times) + 1):
            axes = figure.add_subplot(2, len(times), i)

            x = torch.linspace(-1, 1, 1000)
            t = torch.ones_like(x) * time

            u = model(t, x)

            axes.plot(torch.detach(x), torch.detach(u))

            axes.set_title(f"$t={time}$")
            axes.set_xlabel("$x$")
            axes.set_ylabel("$u(t, x)$")

    _plot_times()

    figure.subplots_adjust(wspace=0.357, hspace=0.502)
    return figure


def main():
    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    checkpoint = checkpoint_manager.last_existed()
    assert checkpoint is not None

    figure = plot_checkpoint(checkpoint[1])
    figure.show()
    matplotlib.pyplot.pause(0)


if __name__ == "__main__":
    main()
